Sort the matrix diagonally [Sort]

Time: O(MxNxLog(min(M,N)); Space: O(MxN); medium

Given a m * n matrix mat of integers, sort it diagonally in ascending order from the top-left to the bottom-right then return the sorted array.

Example 1:

Input: mat = [[3,3,1,1],[2,2,1,2],[1,1,1,2]]

Output: [[1,1,1,1],[1,2,2,2],[1,2,3,3]]

Notes:

  • len(mat) == m

  • len(mat[i]) == n

  • 1 <= m, n <= 100

  • 1 <= mat[i][j] <= 100

Hints:

  1. Use a data structure to store all values of each diagonal.

  2. How to index the data structure with the id of the diagonal?

  3. All cells in the same diagonal (i,j) have the same difference so we can get the diagonal of a cell using the difference i-j.

[3]:

import collections

class Solution1(object):
    def diagonalSort(self, mat):
        """
        :type mat: List[List[int]]
        :rtype: List[List[int]]
        """
        lookup = collections.defaultdict(list)

        for i in range(len(mat)):
            for j in range(len(mat[0])):
                lookup[i-j].append(mat[i][j])

        for v in lookup.values():
            v.sort()

        for i in reversed(range(len(mat))):
            for j in reversed(range(len(mat[0]))):
                mat[i][j] = lookup[i-j].pop()

        return mat

[4]:

s = Solution1()
mat = [[3,3,1,1],[2,2,1,2],[1,1,1,2]]
assert s.diagonalSort(mat) == [[1,1,1,1],[1,2,2,2],[1,2,3,3]]